"""
WebSocket实时通信
"""
import json
import asyncio
from typing import Dict, Set
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
from loguru import logger

from app.services.tracking_service import LocationService
from app.utils.dependencies import get_current_user
from app.models.user import User

router = APIRouter(prefix="/ws", tags=["WebSocket"])


class ConnectionManager:
    """WebSocket连接管理器"""
    
    def __init__(self):
        # 存储活跃连接
        self.active_connections: Dict[str, WebSocket] = {}
        # 存储订阅的设备
        self.device_subscriptions: Dict[str, Set[str]] = {}
        # 存储用户连接
        self.user_connections: Dict[int, Set[str]] = {}
    
    async def connect(self, websocket: WebSocket, client_id: str, user_id: int = None):
        """接受WebSocket连接"""
        await websocket.accept()
        self.active_connections[client_id] = websocket
        
        if user_id:
            if user_id not in self.user_connections:
                self.user_connections[user_id] = set()
            self.user_connections[user_id].add(client_id)
        
        logger.info(f"WebSocket连接建立: {client_id}")
    
    def disconnect(self, client_id: str, user_id: int = None):
        """断开WebSocket连接"""
        if client_id in self.active_connections:
            del self.active_connections[client_id]
        
        if client_id in self.device_subscriptions:
            del self.device_subscriptions[client_id]
        
        if user_id and user_id in self.user_connections:
            self.user_connections[user_id].discard(client_id)
            if not self.user_connections[user_id]:
                del self.user_connections[user_id]
        
        logger.info(f"WebSocket连接断开: {client_id}")
    
    async def send_personal_message(self, message: str, client_id: str):
        """发送个人消息"""
        if client_id in self.active_connections:
            try:
                await self.active_connections[client_id].send_text(message)
            except Exception as e:
                logger.error(f"发送消息失败: {e}")
                self.disconnect(client_id)
    
    async def send_to_user(self, message: str, user_id: int):
        """发送消息给特定用户的所有连接"""
        if user_id in self.user_connections:
            for client_id in self.user_connections[user_id].copy():
                await self.send_personal_message(message, client_id)
    
    async def broadcast(self, message: str):
        """广播消息给所有连接"""
        for client_id in list(self.active_connections.keys()):
            await self.send_personal_message(message, client_id)
    
    async def send_to_device_subscribers(self, device_id: str, message: str):
        """发送消息给订阅特定设备的客户端"""
        for client_id, subscribed_devices in self.device_subscriptions.items():
            if device_id in subscribed_devices:
                await self.send_personal_message(message, client_id)
    
    def subscribe_device(self, client_id: str, device_id: str):
        """订阅设备更新"""
        if client_id not in self.device_subscriptions:
            self.device_subscriptions[client_id] = set()
        self.device_subscriptions[client_id].add(device_id)
        logger.info(f"客户端 {client_id} 订阅设备 {device_id}")
    
    def unsubscribe_device(self, client_id: str, device_id: str):
        """取消订阅设备更新"""
        if client_id in self.device_subscriptions:
            self.device_subscriptions[client_id].discard(device_id)
            logger.info(f"客户端 {client_id} 取消订阅设备 {device_id}")


# 全局连接管理器
manager = ConnectionManager()


@router.websocket("/tracking/{client_id}")
async def websocket_tracking_endpoint(websocket: WebSocket, client_id: str):
    """WebSocket追踪端点"""
    await manager.connect(websocket, client_id)
    
    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_text()
            
            try:
                message = json.loads(data)
                message_type = message.get("type")
                
                if message_type == "subscribe_device":
                    # 订阅设备更新
                    device_id = message.get("device_id")
                    if device_id:
                        manager.subscribe_device(client_id, device_id)
                        await manager.send_personal_message(
                            json.dumps({
                                "type": "subscription_confirmed",
                                "device_id": device_id,
                                "message": f"已订阅设备 {device_id} 的更新"
                            }),
                            client_id
                        )
                
                elif message_type == "unsubscribe_device":
                    # 取消订阅设备更新
                    device_id = message.get("device_id")
                    if device_id:
                        manager.unsubscribe_device(client_id, device_id)
                        await manager.send_personal_message(
                            json.dumps({
                                "type": "unsubscription_confirmed",
                                "device_id": device_id,
                                "message": f"已取消订阅设备 {device_id} 的更新"
                            }),
                            client_id
                        )
                
                elif message_type == "get_realtime_locations":
                    # 获取实时位置
                    try:
                        locations = await LocationService.get_latest_locations()
                        await manager.send_personal_message(
                            json.dumps({
                                "type": "realtime_locations",
                                "data": [location.dict() for location in locations]
                            }),
                            client_id
                        )
                    except Exception as e:
                        await manager.send_personal_message(
                            json.dumps({
                                "type": "error",
                                "message": f"获取实时位置失败: {str(e)}"
                            }),
                            client_id
                        )
                
                elif message_type == "ping":
                    # 心跳检测
                    await manager.send_personal_message(
                        json.dumps({
                            "type": "pong",
                            "timestamp": message.get("timestamp")
                        }),
                        client_id
                    )
                
                else:
                    await manager.send_personal_message(
                        json.dumps({
                            "type": "error",
                            "message": f"未知消息类型: {message_type}"
                        }),
                        client_id
                    )
                    
            except json.JSONDecodeError:
                await manager.send_personal_message(
                    json.dumps({
                        "type": "error",
                        "message": "无效的JSON格式"
                    }),
                    client_id
                )
            except Exception as e:
                logger.error(f"处理WebSocket消息失败: {e}")
                await manager.send_personal_message(
                    json.dumps({
                        "type": "error",
                        "message": f"处理消息失败: {str(e)}"
                    }),
                    client_id
                )
                
    except WebSocketDisconnect:
        manager.disconnect(client_id)
    except Exception as e:
        logger.error(f"WebSocket连接异常: {e}")
        manager.disconnect(client_id)


async def broadcast_location_update(device_id: str, location_data: dict):
    """广播位置更新"""
    message = json.dumps({
        "type": "location_update",
        "device_id": device_id,
        "data": location_data
    })
    
    await manager.send_to_device_subscribers(device_id, message)


async def broadcast_geofence_event(event_data: dict):
    """广播地理围栏事件"""
    message = json.dumps({
        "type": "geofence_event",
        "data": event_data
    })
    
    await manager.broadcast(message)


async def broadcast_device_status_change(device_id: str, status: str):
    """广播设备状态变化"""
    message = json.dumps({
        "type": "device_status_change",
        "device_id": device_id,
        "status": status
    })
    
    await manager.send_to_device_subscribers(device_id, message)


# 定时任务：定期广播实时位置
async def periodic_location_broadcast():
    """定期广播实时位置"""
    while True:
        try:
            if manager.active_connections:
                locations = await LocationService.get_latest_locations()
                message = json.dumps({
                    "type": "periodic_location_update",
                    "data": [location.dict() for location in locations]
                })
                await manager.broadcast(message)
        except Exception as e:
            logger.error(f"定期位置广播失败: {e}")
        
        # 每30秒广播一次
        await asyncio.sleep(30)


# 启动定时任务
async def start_background_tasks():
    """启动后台任务"""
    asyncio.create_task(periodic_location_broadcast())
